from concurrent.futures import ThreadPoolExecutor

import pika
import json
import time
import requests  # 实际调用大模型接口用

import mq


class RPCServer:
    def __init__(self):
        self.connection = pika.BlockingConnection(
            pika.ConnectionParameters(host='192.168.122.173'))
        self.channel = self.connection.channel()
        self.channel.queue_declare(queue='rpc_queue', durable=True)
        self.channel.basic_qos(prefetch_count=50)
        self.executor = ThreadPoolExecutor(max_workers=50)  # 创建一个包含10个工作线程的线程池

    def on_request(self, ch, method, props, body):
        try:
            # 解析请求数据
            request_data = json.loads(body)
            print(f" [x] 收到请求: {request_data['id']} ")
            # 调用大模型接口（示例使用模拟数据）
            response = self.call_llm_api(request_data)
            # 发送响应
            ch.basic_publish(
                exchange='',
                routing_key=props.reply_to,
                properties=pika.BasicProperties(
                    correlation_id=props.correlation_id,
                    delivery_mode=2
                ),
                body=json.dumps(response)
            )
            ch.basic_ack(delivery_tag=method.delivery_tag)

        except Exception as e:
            print(f"处理失败: {str(e)}")
            ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)

    def call_llm_api(self, data):
        """调用大模型接口（实际需替换为真实API调用）"""
        # 模拟处理延迟
        print(json.dumps(data))
        resp = requests.post(url="http://localhost:8000/v1/chat/completions", json=json.dumps(data))
        return resp.json()

    def start(self):
        self.channel.basic_consume(
            queue='rpc_queue',
            on_message_callback=lambda ch, method, props, body: self.executor.submit(self.on_request, ch, method, props,
                                                                                     body)
        )
        print(" [*] 等待RPC请求，按 CTRL+C 退出")
        self.channel.start_consuming()


if __name__ == "__main__":
    server = RPCServer()
    server.start()
